/**
 * @file pose_detect.cc
 * @author your name (you@domain.com)
 * @brief 人体关键点检测相关事务
 * @version 0.1
 * @date 2022-09-17
 * 
 * @copyright Copyright (c) 2022
 * 
 */
#include "imageProcess.h"
#include "face_detect.h"
#include "object_detect.h"          //需要object_detect.c中定义的load_model()函数
#include "rknn_api.h"
#include <math.h>
#include "pose_detect.h"
#include <semaphore.h> 

extern queue<pair<vector<int>, Mat>> queueOutput_368;
extern char *pose_model_path;

static rknn_context ctx;
static rknn_input inputs[1];
static rknn_output outputs[1];
 

static int init_rknn_model(){
    int ret;
    int model_len = 0;
    unsigned char *model;
    model = load_model(pose_model_path, &model_len);
    ret = rknn_init(&ctx, model, model_len, 0);
    if(ret < 0) {
        printf("rknn_init fail! ret=%d\n", ret);
        return -1;
    }
    return ret;
}

static void init_rknn_inoutput(){
    // Set Input Data
    inputs[0].index = 0;
    inputs[0].type = RKNN_TENSOR_UINT8;
    inputs[0].size = 368 * 368 * 3;             //width * height * channel;
    inputs[0].fmt = RKNN_TENSOR_NHWC;
    outputs[0].want_float = 1;
}

inline int clamp(float val, int min, int max)
{
    return val > min ? (val < max ? val : max) : min;
}

// @brief 输入一张待检测图，通过数组的形式传回
// @param img 
// @return ret为0 表示检测没有错误
static int detect(Mat img, uint16_t *results){
    int ret; 
    inputs[0].buf = img.data;
    ret = rknn_inputs_set(ctx, 1, inputs);
    ret = rknn_run(ctx, NULL);
    ret = rknn_outputs_get(ctx, 1, outputs, NULL);
    if(ret != 0){
        printf("pose推理出现问题\n");
        return -1;
    }
    float *buffer = (float *)outputs[0].buf;
    float scale = (float)368 / 640;
    for(int i = 0;i<18;i++)
    {   
        int x_pos = 0;
        int y_pos = 0;
        float tempMaxProb = 0;
        //输出向量顺序：沿列的方向进行存储
        for(int j =0;j<46;j++)
        {
            for(int k =0;k<46;k++)
            {
                int offset = i * 46 * 46 + j * 46 + k;
                //是否大于给定阈值
                if(buffer[offset] > 0.1)
                {   
                    if(buffer[offset] > tempMaxProb)
                    {
                        x_pos = k;
                        y_pos = j;
                        tempMaxProb = buffer[offset];
                    }
                }
            }
        }
        if(tempMaxProb > 0.1)
        {
            //将预测坐标映射回原图
            int map_x = (int)(8 * x_pos / scale);
            int map_y = (int)((8 * y_pos) / scale);
            map_x = clamp(map_x,0,640);  
            map_y = clamp(map_y,0,640); 
            //nums数组从0到35的位置依次存放的是每个关键点的xy坐标
            results[2*i] = map_x;
            results[2*i + 1] =  map_y;
            //printf("关键点索引%d,x方向%d,y方向%d\n",i,map_x,map_y);
        }
    }
    return ret;
}
// @brief 该函数用于针对一张图， 判断两个模型的输出是否符合分心驾驶的条件, 符合返回0，否则返回-1
// @param obj_pos 
// @param nums 
static int judge(vector<int> obj_pos, uint16_t *nums)
{
    /*Todo:判定法则 - 两个手腕（如果有）到目标框中心的距离，鼻子（如果有）到目标中心的距离小于阈值返回0。都没有或者距离大于阈值返回-1*/
    int mid_x = (obj_pos[5] + obj_pos[7]) * 0.5 ;
    int mid_y = (obj_pos[6] + obj_pos[8]) * 0.5 ;
    
    int dis = 0;

    //右手腕 - 4 右手肘 -3  左手腕 - 7 左手肘 -6
    // if(nums[0] != 0)
    // {
    //     dis = pow((nums[0] - mid_x),2) +  pow((nums[1] - mid_y),2);
    //     cout << "dis:" << dis << endl;
    //     if(dis <= THRESH)
    //         return 0;
    // }
    // return -1;
    return 0;
}
// @brief 检测送入图片的人体关键点, 并判断和目标框中心的距离,如果确认为手持物体,返回0, 否则返回-1
// @return 0-手持；-1-没有手持
int pose_detect(){ 
    int ret;
    ret = init_rknn_model();
    if(ret != 0)             return -1;
    init_rknn_inoutput();
    
    //主逻辑
    Mat img_368;
    uint16_t nums[36];

    int cnt = 0;
    int n = queueOutput_368.size();
    while(n--)
    {
        pair<vector<int>, Mat> img_info = queueOutput_368.front();
        queueOutput_368.pop();
        cv::resize(img_info.second, img_368, cv::Size(368, 368), cv::INTER_LINEAR);
        memset(nums, 0, sizeof(nums));
        ret = detect(img_368, nums);
        if(ret <= -1){
            cout << "后处理异常，退出！"<<endl;
            return -1;
        }
        //利用nums数据和目标检测框进行综合判断
        vector<int> object_pos = img_info.first;
        
        ret = judge(object_pos, nums);
        if(!ret)     cnt++;

        if(cnt >= 3){
            //说明确实是分心驾驶
            return 0;
        }
    }
    return -1;
}